from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

def simple_query():
    # 初始化嵌入模型
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    
    # 加载向量数据库
    vectorstore = Chroma(
        persist_directory="./chroma_db",
        embedding_function=embeddings
    )
    
    while True:
        query = input("\n请输入查询内容(输入'quit'退出): ").strip()
        
        if query.lower() == 'quit':
            break
            
        if not query:
            print("查询内容不能为空")
            continue
        
        # 执行相似度搜索
        results = vectorstore.similarity_search(query, k=3)
        
        print(f"\n找到 {len(results)} 个相关结果:")
        for i, doc in enumerate(results):
            print(f"\n--- 结果 {i+1} ---")
            print(f"内容: {doc.page_content[:300]}...")  # 限制显示长度
            if hasattr(doc, 'metadata'):
                print(f"来源: {doc.metadata.get('source', '未知')}")
                print(f"页码: {doc.metadata.get('page', '未知')}")

if __name__ == "__main__":
    simple_query()